import albumentations as albu

import cv2


def base_aug():
    return albu.Compose([
        albu.HorizontalFlip(p=0.5),
        albu.VerticalFlip(p=0.5),
        albu.RandomRotate90(p=0.5)
    ])


def aug_1():
    return albu.Compose([
        albu.ShiftScaleRotate(shift_limit=0.2,
                              scale_limit=0.2,
                              rotate_limit=120,
                              border_mode=cv2.BORDER_CONSTANT,
                              value=0, mask_value=0,
                              p=0.5),
        albu.HorizontalFlip(p=0.5),
        albu.VerticalFlip(p=0.5),
    ])
